#!/usr/bin/env python3

"""
Module for backing up and restoring a SpecDB database
"""

SQLITE_PAGE_SIZE_INDEX  = 16
SQLITE_HEADER_LENGTH    = 16
SQLITE_PAGE_COUNT_INDEX = 28

import hashlib
from hashlib import sha256
import os
import sqlite3
import sys

def backup(db=None, object_dir='objects', backup_file='backup.txt'):
	"""
	This function performs the incremental backup 
	This function is taken with slight modifications from the following Github
	repository: https://github.com/nokibsarkar/sqlite3-incremental-backup.git
	
	For this function, we implemented the python version of the backup at 
	sqlite3-incremental-backup/python/sqlite3backup/python 
	
	Much credit to Github user nokibsarkar
	
	Parameters
	----------
	+ db				path to sqlite database to backup
	+ object_dir		path to objects directory where all pages will reside
	+ backup_file		file to save sha256 hashes of pages
	
	Returns
	-------
	True	if backup successful
	"""
	
	page_size = 0
	# Open the database.
	with open(db, "rb") as db_file_object:
		assert(
			db_file_object.read(SQLITE_HEADER_LENGTH) == b"SQLite format 3\x00")
		db_file_object.seek(SQLITE_PAGE_SIZE_INDEX, os.SEEK_SET)
		page_size = int.from_bytes(db_file_object.read(2), 'little') * 256
		db_file_object.seek(SQLITE_PAGE_COUNT_INDEX, os.SEEK_SET)
		page_count = int.from_bytes(db_file_object.read(4), 'big')
	
	pages = []
	with open(db, "rb") as db_file_object:
		for page_number in range(page_count):
			db_file_object.seek(page_number * page_size, os.SEEK_SET)
			page = db_file_object.read(page_size)
			hash = sha256(page).hexdigest()
			directory, filename = hash[:2], hash[2:]
			file_path = os.path.join(object_dir, directory, filename)
			if not os.path.exists(file_path): # 
				os.makedirs(os.path.dirname(file_path), exist_ok=True)
				with open(file_path, "wb") as file_object:
					file_object.write(page)
			pages.append(hash)
	
	# Write the pages to the object directory.
	with open(backup_file, 'w') as fp:
		fp.write('\n'.join(pages))

	return True


def restore(backup=None, backup_file=None, object_dir=None):
	"""
	This function performs the restore function from an incremental backup 
	This function is taken with slight modifications from the following Github
	repository: https://github.com/nokibsarkar/sqlite3-incremental-backup.git
	
	For this function, we implemented the python version of the backup at 
	sqlite3-incremental-backup/python/sqlite3backup/python 
	
	Much credit to Github user nokibsarkar
	
	Parameters
	----------
	+ backup			path to sqlite database to restore to
	+ object_dir		path to objects directory where all pages will reside
	+ backup_file 		file to save sha256 hashes of pages
	
	Returns
	-------
	True	if restore successful
	"""
	
	# Read the pages from the backup file
	with open(backup_file, 'r') as fp:
		pages = fp.read().split('\n')
	
	# Open the database.
	with open(backup, "wb") as db_file_object:
		# Iterate thourgh the pages and write them to the database.
		for page in pages:
			path = os.path.join(object_dir, page[:2], page[2:])
			with open(path, "rb") as file_object:
				db_file_object.write(file_object.read())
	
	# Restoration is complete
	return True
